// Copyright Epic Games, Inc. All Rights Reserved.

#include "zenstore/workspaces.h"

#include <zencore/compactbinarybuilder.h>
#include <zencore/fmtutils.h>
#include <zencore/scopeguard.h>
#include <zencore/timer.h>
#include <zencore/trace.h>
#include <zencore/workthreadpool.h>
#include <zenutil/basicfile.h>

#if ZEN_WITH_TESTS
#	include <zencore/blake3.h>
#	include <zencore/testing.h>
#	include <zencore/testutils.h>
#endif

namespace zen {

namespace {

	std::string WorkspacesToJson(std::span<Workspaces::WorkspaceInfo> Workspaces)
	{
		using namespace std::literals;

		CbObjectWriter Writer;

		Writer.BeginArray("workspaces");

		for (const Workspaces::WorkspaceInfo& Workspace : Workspaces)
		{
			Writer.BeginObject();
			{
				Writer.AddObjectId("id"sv, Workspace.Config.Id);
				Writer.AddString("root_path"sv, reinterpret_cast<const char*>(Workspace.Config.RootPath.u8string().c_str()));
				Writer.BeginArray("shares");
				for (const Workspaces::WorkspaceShareConfiguration& Share : Workspace.Shares)
				{
					Writer.BeginObject();
					{
						Writer.AddObjectId("id"sv, Share.Id);
						Writer.AddString("share_path"sv, reinterpret_cast<const char*>(Share.SharePath.u8string().c_str()));
						if (!Share.Alias.empty())
						{
							Writer.AddString("alias"sv, Share.Alias);
						}
					}
					Writer.EndObject();
				}
				Writer.EndArray();
			}
			Writer.EndObject();
		}
		Writer.EndArray();
		ExtendableStringBuilder<512> Json;
		Writer.Save().ToJson(Json);
		return Json.ToString();
	}

	std::vector<Workspaces::WorkspaceInfo> WorkspacesFromJson(const IoBuffer&											   WorkspaceJson,
															  const std::function<Oid(const std::filesystem::path& Path)>& PathToIdCB,
															  std::string&												   OutError)
	{
		using namespace std::literals;

		CbFieldIterator RootField =
			LoadCompactBinaryFromJson(std::string_view((const char*)(WorkspaceJson.Data()), WorkspaceJson.GetSize()), OutError);
		if (OutError.empty())
		{
			std::vector<Workspaces::WorkspaceInfo> Workspaces;
			if (CbObjectView RootObject = RootField.AsObjectView(); RootObject)
			{
				for (CbFieldView WorkspaceField : RootObject["workspaces"].AsArrayView())
				{
					CbObjectView		  Workspace	  = WorkspaceField.AsObjectView();
					Oid					  WorkspaceId = Workspace["id"sv].AsObjectId();
					std::filesystem::path RootPath	  = Workspace["root_path"sv].AsU8String();
					if (WorkspaceId == Oid::Zero && !RootPath.empty())
					{
						WorkspaceId = PathToIdCB(RootPath);
					}
					if (WorkspaceId != Oid::Zero && !RootPath.empty())
					{
						std::vector<Workspaces::WorkspaceShareConfiguration> Shares;
						for (CbFieldView ShareField : Workspace["shares"].AsArrayView())
						{
							CbObjectView		  Share		= ShareField.AsObjectView();
							Oid					  ShareId	= Share["id"sv].AsObjectId();
							std::filesystem::path SharePath = Share["share_path"sv].AsU8String();
							if (ShareId == Oid::Zero && !SharePath.empty())
							{
								ShareId = PathToIdCB(SharePath);
							}
							std::string_view Alias = Share["alias"sv].AsString();
							if (ShareId != Oid::Zero && !SharePath.empty())
							{
								Shares.push_back({.Id = ShareId, .SharePath = SharePath, .Alias = std::string(Alias)});
							}
						}
						Workspaces.push_back({.Config = {.Id = WorkspaceId, .RootPath = RootPath}, .Shares = std::move(Shares)});
					}
				}
			}
			return Workspaces;
		}
		return {};
	}

}  // namespace
//////////////////////////////////////////////////////////////////////////

class FolderStructure
{
public:
	struct FileEntry
	{
		std::filesystem::path RelativePath;
		uint64_t			  Size;
	};

	FolderStructure() {}
	FolderStructure(std::vector<FileEntry>&& InEntries, std::vector<Oid>&& Ids);

	const FileEntry* FindEntry(const Oid& Id) const
	{
		if (auto It = IdLookup.find(Id); It != IdLookup.end())
		{
			return &Entries[It->second];
		}
		return nullptr;
	}

	size_t EntryCount() const { return Entries.size(); }

	void IterateEntries(std::function<void(const Oid& Id, const FileEntry& Entry)>&& Callback) const
	{
		for (auto It = IdLookup.begin(); It != IdLookup.end(); It++)
		{
			Callback(It->first, Entries[It->second]);
		}
	}

private:
	const std::vector<FileEntry>			 Entries;
	tsl::robin_map<Oid, size_t, Oid::Hasher> IdLookup;
};

//////////////////////////////////////////////////////////////////////////

class WorkspaceShare : public RefCounted
{
public:
	WorkspaceShare(const Workspaces::WorkspaceShareConfiguration&				Config,
				   std::unique_ptr<FolderStructure>&&							FolderStructure,
				   const std::function<Oid(const std::filesystem::path& Path)>& PathToId);

	const Workspaces::WorkspaceShareConfiguration& GetConfig() const;

	bool IsInitialized() const { return !!m_FolderStructure; }

	const std::function<Oid(const std::filesystem::path& Path)>& GetPathToIdFunction() const { return m_PathToid; }

	std::filesystem::path GetAbsolutePath(const std::filesystem::path& RootPath, const Oid& ChunkId, uint64_t& OutSize) const;

	const FolderStructure& GetStructure() const
	{
		ZEN_ASSERT(m_FolderStructure);
		return *m_FolderStructure;
	}

private:
	const Workspaces::WorkspaceShareConfiguration		  m_Config;
	std::function<Oid(const std::filesystem::path& Path)> m_PathToid;
	std::unique_ptr<FolderStructure>					  m_FolderStructure;
};

//////////////////////////////////////////////////////////////////////////

class Workspace : public RefCounted
{
public:
	Workspace(LoggerRef& Log, const Workspaces::WorkspaceConfiguration& Config);

	const Workspaces::WorkspaceConfiguration& GetConfig() const;
	std::vector<Ref<WorkspaceShare>>		  GetShares() const;
	Ref<WorkspaceShare>						  GetShare(const Oid& ShareId) const;

	void SetShare(const Oid& ShareId, Ref<WorkspaceShare>&& Share);

private:
	LoggerRef Log() { return m_Log; }

	LoggerRef&											  m_Log;
	const Workspaces::WorkspaceConfiguration			  m_Config;
	tsl::robin_map<Oid, Ref<WorkspaceShare>, Oid::Hasher> m_Shares;
};

//////////////////////////////////////////////////////////////////////////

FolderStructure::FolderStructure(std::vector<FileEntry>&& InEntries, std::vector<Oid>&& Ids) : Entries(std::move(InEntries))
{
	IdLookup.reserve(Entries.size());
	for (size_t Index = 0; Index < Entries.size(); Index++)
	{
		Oid Id = Ids[Index];
		IdLookup.insert(std::make_pair(Id, Index));
	}
}

namespace {
	struct FolderScanner
	{
		FolderScanner(LoggerRef&												   Log,
					  WorkerThreadPool&											   WorkerPool,
					  const std::filesystem::path&								   Path,
					  const std::function<Oid(const std::filesystem::path& Path)>& PathToIdCB)
		: m_Log(Log)
		, Path(Path)
		, PathToIdCB(PathToIdCB)
		, WorkLatch(1)
		, WorkerPool(WorkerPool)
		{
		}

		void Traverse();
		void Traverse(const std::filesystem::path& RelativeRoot, const std::filesystem::path& Path);

		LoggerRef&													 Log() { return m_Log; }
		LoggerRef&													 m_Log;
		const std::filesystem::path									 Path;
		RwLock														 WorkLock;
		const std::function<Oid(const std::filesystem::path& Path)>& PathToIdCB;
		std::vector<FolderStructure::FileEntry>						 FoundFiles;
		std::vector<Oid>											 FoundFileIds;
		Latch														 WorkLatch;
		WorkerThreadPool&											 WorkerPool;
	};

	struct Visitor : public FileSystemTraversal::TreeVisitor
	{
		Visitor(FolderScanner& Data, const std::filesystem::path& RelativeRoot) : Data(Data), RelativeRoot(RelativeRoot) {}

		FileSystemTraversal						Traverser;
		FolderScanner&							Data;
		std::vector<FolderStructure::FileEntry> Entries;
		std::vector<Oid>						FileIds;
		std::filesystem::path					RelativeRoot;

		virtual void VisitFile(const std::filesystem::path&, const path_view& File, uint64_t FileSize)
		{
			std::filesystem::path RelativePath = RelativeRoot.empty() ? File : RelativeRoot / File;
			Entries.push_back(FolderStructure::FileEntry{.RelativePath = RelativePath, .Size = FileSize});
			FileIds.push_back(Data.PathToIdCB(RelativePath));
		}

		virtual bool VisitDirectory(const std::filesystem::path& Parent, const path_view& DirectoryName)
		{
			ZEN_ASSERT(!Parent.empty());
			ZEN_ASSERT(!DirectoryName.empty());
			FolderScanner* DataPtr = &Data;
			Data.WorkLatch.AddCount(1);
			Data.WorkerPool.ScheduleWork([DataPtr,
										  RootDir	   = Parent / DirectoryName,
										  RelativeRoot = RelativeRoot.empty() ? DirectoryName : RelativeRoot / DirectoryName]() {
				auto _ = MakeGuard([DataPtr]() { DataPtr->WorkLatch.CountDown(); });
				DataPtr->Traverse(RelativeRoot, RootDir);
			});
			return false;
		}
	};

	void FolderScanner::Traverse()
	{
		Stopwatch Timer;
		Traverse({}, std::filesystem::absolute(Path));
		WorkLatch.CountDown();
		while (!WorkLatch.Wait(1000))
		{
			WorkLock.WithSharedLock([&]() { ZEN_INFO("Found {} files in '{}'...", FoundFiles.size(), Path.string()); });
		}
		ZEN_ASSERT(FoundFiles.size() == FoundFileIds.size());
		ZEN_INFO("Found {} files in '{}' in {}", FoundFiles.size(), Path.string(), NiceLatencyNs(Timer.GetElapsedTimeUs() * 1000));
	}

	void FolderScanner::Traverse(const std::filesystem::path& RelativeRoot, const std::filesystem::path& AbsoluteRoot)
	{
		Visitor LeafVisitor(*this, RelativeRoot);
		LeafVisitor.Traverser.TraverseFileSystem(AbsoluteRoot, LeafVisitor);
		if (!LeafVisitor.Entries.empty())
		{
			WorkLock.WithExclusiveLock([&]() {
				FoundFiles.insert(FoundFiles.end(), LeafVisitor.Entries.begin(), LeafVisitor.Entries.end());
				FoundFileIds.insert(FoundFileIds.end(), LeafVisitor.FileIds.begin(), LeafVisitor.FileIds.end());
			});
		}
	}
}  // namespace

std::unique_ptr<FolderStructure>
ScanFolder(LoggerRef													InLog,
		   const std::filesystem::path&									Path,
		   const std::function<Oid(const std::filesystem::path& Path)>& PathToIdCB,
		   WorkerThreadPool&											WorkerPool)
{
	ZEN_TRACE_CPU("workspaces::ScanFolderImpl");

	auto Log = [&InLog]() { return InLog; };

	FolderScanner Data(InLog, WorkerPool, Path, PathToIdCB);
	Data.Traverse();
	return std::make_unique<FolderStructure>(std::move(Data.FoundFiles), std::move(Data.FoundFileIds));
}

////////////////////////////////////////////////////////////

WorkspaceShare::WorkspaceShare(const Workspaces::WorkspaceShareConfiguration&				Config,
							   std::unique_ptr<FolderStructure>&&							FolderStructure,
							   const std::function<Oid(const std::filesystem::path& Path)>& PathToId)
: m_Config(Config)
, m_PathToid(PathToId)
, m_FolderStructure(std::move(FolderStructure))
{
}

std::filesystem::path
WorkspaceShare::GetAbsolutePath(const std::filesystem::path& RootPath, const Oid& Id, uint64_t& OutSize) const
{
	ZEN_ASSERT(m_FolderStructure);
	const FolderStructure::FileEntry* Entry = m_FolderStructure->FindEntry(Id);
	if (Entry == nullptr)
	{
		return {};
	}
	OutSize = Entry->Size;
	return RootPath / m_Config.SharePath / Entry->RelativePath;
}

const Workspaces::WorkspaceShareConfiguration&
WorkspaceShare::GetConfig() const
{
	return m_Config;
}

////////////////////////////////////////////////////////////

Workspace::Workspace(LoggerRef& Log, const Workspaces::WorkspaceConfiguration& Config) : m_Log(Log), m_Config(Config)
{
}

const Workspaces::WorkspaceConfiguration&
Workspace::GetConfig() const
{
	return m_Config;
}
std::vector<Ref<WorkspaceShare>>
Workspace::GetShares() const
{
	std::vector<Ref<WorkspaceShare>> Shares;
	Shares.reserve(m_Shares.size());
	for (auto It : m_Shares)
	{
		Shares.push_back(It.second);
	}
	return Shares;
}

Ref<WorkspaceShare>
Workspace::GetShare(const Oid& ShareId) const
{
	if (auto It = m_Shares.find(ShareId); It != m_Shares.end())
	{
		return It->second;
	}
	return {};
}

void
Workspace::SetShare(const Oid& ShareId, Ref<WorkspaceShare>&& Share)
{
	if (Share)
	{
		m_Shares.insert_or_assign(ShareId, std::move(Share));
	}
	else
	{
		m_Shares.erase(ShareId);
	}
}

////////////////////////////////////////////////////////////

Workspaces::Workspaces() : m_Log(logging::Get("workspaces"))
{
}

Workspaces::~Workspaces()
{
}

bool
Workspaces::AddWorkspace(const WorkspaceConfiguration& Configuration)
{
	Ref<Workspace> NewWorkspace(new Workspace(m_Log, Configuration));

	RwLock::ExclusiveLockScope Lock(m_Lock);
	if (m_Workspaces.contains(Configuration.Id))
	{
		return false;
	}

	const std::filesystem::path RootPath = Configuration.RootPath;
	if (RootPath.is_relative())
	{
		throw std::invalid_argument(fmt::format("workspace root path '{}' is not an absolute path", RootPath));
	}

	m_Workspaces.insert(std::make_pair(Configuration.Id, NewWorkspace));
	ZEN_INFO("Created workspace '{}' with root '{}'", Configuration.Id, Configuration.RootPath);
	return true;
}

Workspaces::WorkspaceConfiguration
Workspaces::GetWorkspaceConfiguration(const Oid& WorkspaceId) const
{
	RwLock::SharedLockScope Lock(m_Lock);
	Ref<Workspace>			Workspace = FindWorkspace(Lock, WorkspaceId);
	if (Workspace)
	{
		return Workspace->GetConfig();
	}
	return {};
}

Workspaces::WorkspaceInfo
Workspaces::GetWorkspaceInfo(const Oid& WorkspaceId) const
{
	Ref<Workspace>					 Workspace;
	std::vector<Ref<WorkspaceShare>> Shares;
	{
		RwLock::SharedLockScope Lock(m_Lock);
		Workspace = FindWorkspace(Lock, WorkspaceId);
		if (Workspace)
		{
			Shares = Workspace->GetShares();
		}
	}
	if (!Workspace)
	{
		return {};
	}

	WorkspaceInfo Info = {.Config = Workspace->GetConfig()};
	Info.Shares.reserve(Shares.size());
	for (const Ref<WorkspaceShare>& Share : Shares)
	{
		Info.Shares.push_back(Share->GetConfig());
	}
	return Info;
}

bool
Workspaces::RemoveWorkspace(const Oid& WorkspaceId)
{
	RwLock::ExclusiveLockScope Lock(m_Lock);
	if (auto It = m_Workspaces.find(WorkspaceId); It != m_Workspaces.end())
	{
		std::vector<std::string> Aliases;
		for (const auto& AliasIt : m_ShareAliases)
		{
			if (AliasIt.second.WorkspaceId == WorkspaceId)
			{
				Aliases.push_back(AliasIt.first);
			}
		}

		for (const std::string& Alias : Aliases)
		{
			m_ShareAliases.erase(Alias);
		}

		m_Workspaces.erase(It);

		ZEN_INFO("Removed workspace '{}' and {} aliases", WorkspaceId, Aliases.size());
		return true;
	}
	return false;
}

bool
Workspaces::AddWorkspaceShare(const Oid&												   WorkspaceId,
							  const WorkspaceShareConfiguration&						   Configuration,
							  const std::function<Oid(const std::filesystem::path& Path)>& PathToIdCB)
{
	Ref<Workspace> Workspace;
	{
		RwLock::SharedLockScope Lock(m_Lock);
		Workspace = FindWorkspace(Lock, WorkspaceId);
		if (!Workspace)
		{
			return false;
		}
		if (Workspace->GetShare(Configuration.Id))
		{
			return false;
		}
	}

	const std::filesystem::path RootPath  = Workspace->GetConfig().RootPath;
	const std::filesystem::path SharePath = Configuration.SharePath;

	std::filesystem::path FullPath		  = std::filesystem::absolute(RootPath / SharePath);
	std::filesystem::path VerifySharePath = std::filesystem::relative(FullPath, RootPath);
	if (VerifySharePath != Configuration.SharePath || VerifySharePath.string().starts_with(".."))
	{
		throw std::invalid_argument(fmt::format("workspace share path '{}' is not valid for root path '{}'", SharePath, RootPath));
	}

	Ref<WorkspaceShare> NewShare(new WorkspaceShare(Configuration, {}, PathToIdCB));
	{
		RwLock::ExclusiveLockScope _(m_Lock);
		Workspace->SetShare(Configuration.Id, std::move(NewShare));
		if (!Configuration.Alias.empty())
		{
			m_ShareAliases.insert_or_assign(Configuration.Alias, ShareAlias{.WorkspaceId = WorkspaceId, .ShareId = Configuration.Id});
		}
	}

	ZEN_INFO("Added workspace share '{}' in workspace '{}' with path '{}'", Configuration.Id, WorkspaceId, Configuration.SharePath);

	return true;
}

Workspaces::WorkspaceShareConfiguration
Workspaces::GetWorkspaceShareConfiguration(const Oid& WorkspaceId, const Oid& ShareId) const
{
	RwLock::SharedLockScope Lock(m_Lock);
	Ref<Workspace>			Workspace = FindWorkspace(Lock, WorkspaceId);
	if (Workspace)
	{
		Ref<WorkspaceShare> Share = Workspace->GetShare(ShareId);
		if (Share)
		{
			return Share->GetConfig();
		}
	}
	return {};
}

bool
Workspaces::RemoveWorkspaceShare(const Oid& WorkspaceId, const Oid& ShareId)
{
	Ref<Workspace> Workspace;
	{
		RwLock::SharedLockScope Lock(m_Lock);
		Workspace = FindWorkspace(Lock, WorkspaceId);
		if (!Workspace)
		{
			return false;
		}
	}
	RwLock::ExclusiveLockScope _(m_Lock);
	Ref<WorkspaceShare>		   ExistingShare = Workspace->GetShare(ShareId);
	if (!ExistingShare)
	{
		return false;
	}

	std::string Alias = ExistingShare->GetConfig().Alias;
	if (!Alias.empty())
	{
		m_ShareAliases.erase(Alias);
	}

	Workspace->SetShare(ShareId, {});

	ZEN_INFO("Removed workspace share '{}' in workspace '{}'", ShareId, WorkspaceId);
	return true;
}

std::optional<std::vector<Workspaces::ShareFile>>
Workspaces::GetWorkspaceShareFiles(const Oid& WorkspaceId, const Oid& ShareId, bool ForceRefresh, WorkerThreadPool& WorkerPool)
{
	std::pair<Ref<Workspace>, Ref<WorkspaceShare>> WorkspaceAndShare = FindWorkspaceShare(WorkspaceId, ShareId, ForceRefresh, WorkerPool);
	if (!WorkspaceAndShare.second)
	{
		return {};
	}

	const FolderStructure&			   Structure = WorkspaceAndShare.second->GetStructure();
	std::vector<Workspaces::ShareFile> Files;
	Files.reserve(Structure.EntryCount());
	Structure.IterateEntries([&Files](const Oid& Id, const FolderStructure::FileEntry& Entry) {
		std::string GenericPath(reinterpret_cast<const char*>(Entry.RelativePath.generic_u8string().c_str()));
		Files.push_back(ShareFile{.RelativePath = std::move(GenericPath), .Size = Entry.Size, .Id = Id});
	});
	return Files;
}

Workspaces::ShareFile
Workspaces::GetWorkspaceShareChunkInfo(const Oid& WorkspaceId, const Oid& ShareId, const Oid& ChunkId, WorkerThreadPool& WorkerPool)
{
	using namespace std::literals;

	std::pair<Ref<Workspace>, Ref<WorkspaceShare>> WorkspaceAndShare = FindWorkspaceShare(WorkspaceId, ShareId, false, WorkerPool);
	if (!WorkspaceAndShare.second)
	{
		return {};
	}

	const FolderStructure::FileEntry* Entry = WorkspaceAndShare.second->GetStructure().FindEntry(ChunkId);
	if (Entry)
	{
		std::string GenericPath(reinterpret_cast<const char*>(Entry->RelativePath.generic_u8string().c_str()));
		return Workspaces::ShareFile{.RelativePath = std::move(GenericPath), .Size = Entry->Size, .Id = ChunkId};
	}
	return {};
}

std::vector<IoBuffer>
Workspaces::GetWorkspaceShareChunks(const Oid&							WorkspaceId,
									const Oid&							ShareId,
									const std::span<const ChunkRequest> ChunkRequests,
									WorkerThreadPool&					WorkerPool)
{
	if (ChunkRequests.size() == 0)
	{
		return {};
	}

	std::pair<Ref<Workspace>, Ref<WorkspaceShare>> WorkspaceAndShare = FindWorkspaceShare(WorkspaceId, ShareId, false, WorkerPool);
	if (!WorkspaceAndShare.second)
	{
		return {};
	}

	std::filesystem::path RootPath = WorkspaceAndShare.first->GetConfig().RootPath;

	auto GetOne = [this](const std::filesystem::path& RootPath, WorkspaceShare& Share, const ChunkRequest& Request) -> IoBuffer {
		uint64_t			  Size;
		std::filesystem::path Path = Share.GetAbsolutePath(RootPath, Request.ChunkId, Size);
		if (!Path.empty())
		{
			uint64_t RequestedOffset = Request.Offset;
			uint64_t RequestedSize	 = Request.Size;
			if (Request.Offset > 0 || Request.Size < uint64_t(-1))
			{
				if (RequestedOffset > Size)
				{
					RequestedOffset = Size;
				}
				if ((RequestedOffset + RequestedSize) > Size)
				{
					RequestedSize = Size - RequestedOffset;
				}
			}
			return IoBufferBuilder::MakeFromFile(Path, RequestedOffset, RequestedSize);
		}
		return IoBuffer{};
	};

	if (ChunkRequests.size() == 1)
	{
		return std::vector<IoBuffer>({GetOne(RootPath, *WorkspaceAndShare.second, ChunkRequests[0])});
	}

	std::vector<IoBuffer> Chunks;
	Chunks.resize(ChunkRequests.size());

	Latch WorkLatch(1);
	for (size_t Index = 0; Index < ChunkRequests.size(); Index++)
	{
		WorkLatch.AddCount(1);
		WorkerPool.ScheduleWork([&, Index]() {
			auto _		  = MakeGuard([&WorkLatch]() { WorkLatch.CountDown(); });
			Chunks[Index] = GetOne(RootPath, *WorkspaceAndShare.second, ChunkRequests[Index]);
		});
	}
	WorkLatch.CountDown();
	WorkLatch.Wait();

	return Chunks;
}

void
Workspaces::WriteState(const std::filesystem::path& WorkspaceStatePath)
{
	using namespace std::literals;

	ZEN_INFO("Writing workspaces state to {}", WorkspaceStatePath);

	CreateDirectories(WorkspaceStatePath);

	std::vector<WorkspaceInfo> Workspaces;
	{
		RwLock::SharedLockScope _(m_Lock);
		for (auto WorkspaceIt : m_Workspaces)
		{
			std::vector<Workspaces::WorkspaceShareConfiguration> Shares;
			for (auto ShareIt : WorkspaceIt.second->GetShares())
			{
				Shares.push_back(ShareIt->GetConfig());
			}
			Workspaces.push_back({.Config = WorkspaceIt.second->GetConfig(), .Shares = std::move(Shares)});
		}
	}
	std::string ConfigJson = WorkspacesToJson(Workspaces);
	TemporaryFile::SafeWriteFile(WorkspaceStatePath / "config.json", MemoryView(ConfigJson.data(), ConfigJson.size()));
}

void
Workspaces::ReadState(const std::filesystem::path&								   WorkspaceStatePath,
					  const std::function<Oid(const std::filesystem::path& Path)>& PathToIdCB)
{
	using namespace std::literals;

	std::string							   Error;
	std::vector<Workspaces::WorkspaceInfo> Workspaces =
		WorkspacesFromJson(IoBufferBuilder::MakeFromFile(WorkspaceStatePath / "config.json"sv), PathToIdCB, Error);
	if (!Error.empty())
	{
		ZEN_WARN("Failed to read workspace state from {}. Reason: '{}'", WorkspaceStatePath / "config.json"sv, Error);
	}
	else
	{
		for (const Workspaces::WorkspaceInfo& Workspace : Workspaces)
		{
			try
			{
				if (AddWorkspace(Workspace.Config))
				{
					for (const Workspaces::WorkspaceShareConfiguration& Share : Workspace.Shares)
					{
						try
						{
							(void)AddWorkspaceShare(Workspace.Config.Id, Share, PathToIdCB);
						}
						catch (const std::exception& Ex)
						{
							ZEN_WARN("Failed adding workspace share '{}' for workspace '{}'. Reason: '{}'",
									 Workspace.Config.Id,
									 Share.Id,
									 Ex.what());
						}
					}
				}
			}
			catch (const std::exception& Ex)
			{
				ZEN_WARN("Failed adding workspace '{}'. Reason: '{}'", Workspace.Config.Id, Ex.what());
			}
		}
	}
}

std::optional<Workspaces::ShareAlias>
Workspaces::GetShareAlias(std::string_view Alias) const
{
	RwLock::SharedLockScope Lock(m_Lock);
	if (auto It = m_ShareAliases.find(std::string(Alias)); It != m_ShareAliases.end())
	{
		return It->second;
	}
	return {};
}

std::pair<Ref<Workspace>, Ref<WorkspaceShare>>
Workspaces::FindWorkspaceShare(const Oid& WorkspaceId, const Oid& ShareId, bool ForceRefresh, WorkerThreadPool& WorkerPool)
{
	Ref<Workspace>		Workspace;
	Ref<WorkspaceShare> Share;
	{
		RwLock::SharedLockScope Lock(m_Lock);
		Workspace = FindWorkspace(Lock, WorkspaceId);
		if (!Workspace)
		{
			return {};
		}
		Share = Workspace->GetShare(ShareId);
		if (!Share)
		{
			return {};
		}
	}

	if (ForceRefresh || !Share->IsInitialized())
	{
		Workspaces::WorkspaceShareConfiguration				  Config	 = Share->GetConfig();
		std::filesystem::path								  RootPath	 = Workspace->GetConfig().RootPath;
		std::function<Oid(const std::filesystem::path& Path)> PathToIdCB = Share->GetPathToIdFunction();
		std::unique_ptr<FolderStructure> NewStructure = ScanFolder(Log(), RootPath / Config.SharePath, PathToIdCB, WorkerPool);
		if (NewStructure)
		{
			Share = Ref<WorkspaceShare>(new WorkspaceShare(Config, std::move(NewStructure), std::move(PathToIdCB)));
			{
				RwLock::ExclusiveLockScope _(m_Lock);
				Workspace->SetShare(ShareId, Ref<WorkspaceShare>(Share));
			}
		}
		else
		{
			if (!Share->IsInitialized())
			{
				ZEN_WARN("Failed to scan folder {} for share {} in workspace {}, treating it as an empty share",
						 WorkspaceId,
						 ShareId,
						 RootPath / Config.SharePath);
				Share = Ref<WorkspaceShare>(new WorkspaceShare(Config, std::move(NewStructure), std::move(PathToIdCB)));
				{
					RwLock::ExclusiveLockScope _(m_Lock);
					Workspace->SetShare(ShareId, Ref<WorkspaceShare>(Share));
				}
			}
		}
	}
	return {std::move(Workspace), std::move(Share)};
}

Ref<Workspace>
Workspaces::FindWorkspace(const RwLock::SharedLockScope&, const Oid& WorkspaceId) const
{
	if (auto It = m_Workspaces.find(WorkspaceId); It != m_Workspaces.end())
	{
		return It->second;
	}
	return {};
}

//////////////////////////////////////////////////////////////////////////

#if ZEN_WITH_TESTS

namespace {
	Oid PathToId(const std::filesystem::path& Path)
	{
		return Oid::FromMemory(BLAKE3::HashMemory((const void*)Path.string().data(), Path.string().size()).Hash);
	}

	std::vector<std::pair<std::filesystem::path, IoBuffer>> GenerateFolderContent(const std::filesystem::path& RootPath)
	{
		std::vector<std::pair<std::filesystem::path, IoBuffer>> Result;
		Result.push_back(std::make_pair(RootPath / "root_blob_1.bin", CreateRandomBlob(4122)));
		Result.push_back(std::make_pair(RootPath / "root_blob_2.bin", CreateRandomBlob(2122)));

		std::filesystem::path EmptyFolder(RootPath / "empty_folder");

		std::filesystem::path FirstFolder(RootPath / "first_folder");
		std::filesystem::create_directory(FirstFolder);
		Result.push_back(std::make_pair(FirstFolder / "first_folder_blob1.bin", CreateRandomBlob(22)));
		Result.push_back(std::make_pair(FirstFolder / "first_folder_blob2.bin", CreateRandomBlob(122)));

		std::filesystem::path SecondFolder(RootPath / "second_folder");
		std::filesystem::create_directory(SecondFolder);
		Result.push_back(std::make_pair(SecondFolder / "second_folder_blob1.bin", CreateRandomBlob(522)));
		Result.push_back(std::make_pair(SecondFolder / "second_folder_blob2.bin", CreateRandomBlob(122)));
		Result.push_back(std::make_pair(SecondFolder / "second_folder_blob3.bin", CreateRandomBlob(225)));

		std::filesystem::path SecondFolderChild(SecondFolder / "child_in_second");
		std::filesystem::create_directory(SecondFolderChild);
		Result.push_back(std::make_pair(SecondFolderChild / "second_child_folder_blob1.bin", CreateRandomBlob(622)));

		for (const auto& It : Result)
		{
			WriteFile(It.first, It.second);
		}

		return Result;
	}

	std::vector<std::pair<std::filesystem::path, IoBuffer>> GenerateFolderContent2(const std::filesystem::path& RootPath)
	{
		std::vector<std::pair<std::filesystem::path, IoBuffer>> Result;
		Result.push_back(std::make_pair(RootPath / "root_blob_3.bin", CreateRandomBlob(312)));
		std::filesystem::path FirstFolder(RootPath / "first_folder");
		Result.push_back(std::make_pair(FirstFolder / "first_folder_blob3.bin", CreateRandomBlob(722)));
		std::filesystem::path SecondFolder(RootPath / "second_folder");
		std::filesystem::path SecondFolderChild(SecondFolder / "child_in_second");
		Result.push_back(std::make_pair(SecondFolderChild / "second_child_folder_blob2.bin", CreateRandomBlob(962)));
		Result.push_back(std::make_pair(SecondFolderChild / "second_child_folder_blob3.bin", CreateRandomBlob(561)));

		for (const auto& It : Result)
		{
			WriteFile(It.first, It.second);
		}

		return Result;
	}

}  // namespace

TEST_CASE("workspaces.scanfolder")
{
	using namespace std::literals;

	WorkerThreadPool WorkerPool(std::thread::hardware_concurrency());

	ScopedTemporaryDirectory TempDir;
	std::filesystem::path	 RootPath = TempDir.Path();
	(void)GenerateFolderContent(RootPath);

	std::unique_ptr<FolderStructure> Structure = ScanFolder(
		logging::Default(),
		RootPath,
		[](const std::filesystem::path& Path) { return PathToId(Path); },
		WorkerPool);
	CHECK(Structure);

	Structure->IterateEntries([&](const Oid& Id, const FolderStructure::FileEntry& Entry) {
		std::filesystem::path AbsPath = RootPath / Entry.RelativePath;
		CHECK(std::filesystem::is_regular_file(AbsPath));
		CHECK(std::filesystem::file_size(AbsPath) == Entry.Size);
		const FolderStructure::FileEntry* FindEntry = Structure->FindEntry(Id);
		CHECK(FindEntry);
		std::filesystem::path Path = RootPath / FindEntry->RelativePath;
		CHECK(AbsPath == Path);
		CHECK(std::filesystem::file_size(AbsPath) == FindEntry->Size);
	});
}

TEST_CASE("workspace.share.paths")
{
	using namespace std::literals;

	WorkerThreadPool WorkerPool(std::thread::hardware_concurrency());

	ScopedTemporaryDirectory								TempDir;
	std::filesystem::path									RootPath = TempDir.Path();
	std::vector<std::pair<std::filesystem::path, IoBuffer>> Content	 = GenerateFolderContent(RootPath);
	Workspaces												WS;
	CHECK(WS.AddWorkspace({PathToId(RootPath), RootPath}));
	CHECK(WS.AddWorkspaceShare(PathToId(RootPath),
							   {PathToId("second_folder/child_in_second"), "second_folder/child_in_second"},
							   [](const std::filesystem::path& Path) { return PathToId(Path); }));
	CHECK_THROWS(WS.AddWorkspaceShare(PathToId(RootPath),
									  {PathToId("../second_folder"), "../second_folder"},
									  [](const std::filesystem::path& Path) { return PathToId(Path); }));
	CHECK(WS.AddWorkspaceShare(PathToId(RootPath), {PathToId("."), "."}, [](const std::filesystem::path& Path) { return PathToId(Path); }));
	CHECK_THROWS(
		WS.AddWorkspaceShare(PathToId(RootPath), {PathToId(".."), ".."}, [](const std::filesystem::path& Path) { return PathToId(Path); }));
	CHECK_THROWS(WS.AddWorkspaceShare(PathToId(RootPath),
									  {PathToId("second_folder/../second_folder/../.."), "second_folder/../second_folder/../.."},
									  [](const std::filesystem::path& Path) { return PathToId(Path); }));
	CHECK_THROWS(WS.AddWorkspaceShare(PathToId(RootPath), {PathToId(RootPath), RootPath}, [](const std::filesystem::path& Path) {
		return PathToId(Path);
	}));
}

TEST_CASE("workspace.share.basic")
{
	using namespace std::literals;

	WorkerThreadPool WorkerPool(std::thread::hardware_concurrency());

	ScopedTemporaryDirectory								TempDir;
	std::filesystem::path									RootPath = TempDir.Path();
	std::vector<std::pair<std::filesystem::path, IoBuffer>> Content	 = GenerateFolderContent(RootPath);

	Workspaces WS;
	CHECK(WS.AddWorkspace({PathToId(RootPath), RootPath}));
	CHECK(WS.AddWorkspaceShare(PathToId(RootPath), {PathToId("second_folder"), "second_folder"}, [](const std::filesystem::path& Path) {
		return PathToId(Path);
	}));
	std::filesystem::path			   SharePath = RootPath / "second_folder";
	std::vector<std::filesystem::path> Paths	 = {{std::filesystem::relative(Content[4].first, SharePath)},
												{std::filesystem::relative(Content[6].first, SharePath)},
												{std::filesystem::relative(Content[7].first, SharePath)},
												{"the_file_that_is_not_there.txt"}};
	std::vector<IoBuffer>			   Chunks	 = WS.GetWorkspaceShareChunks(PathToId(RootPath),
															  PathToId("second_folder"),
															  std::vector<Workspaces::ChunkRequest>{{.ChunkId = PathToId(Paths[0])},
																									{.ChunkId = PathToId(Paths[1])},
																									{.ChunkId = PathToId(Paths[2])},
																									{.ChunkId = PathToId(Paths[3])}},
															  WorkerPool);
	CHECK(Chunks.size() == 4);
	CHECK(Chunks[0].GetView().EqualBytes(Content[4].second.GetView()));
	CHECK(Chunks[1].GetView().EqualBytes(Content[6].second.GetView()));
	CHECK(Chunks[2].GetView().EqualBytes(Content[7].second.GetView()));
	CHECK(Chunks[3].GetSize() == 0);

	std::vector<std::pair<std::filesystem::path, IoBuffer>> Content2 = GenerateFolderContent2(RootPath);
	std::vector<std::filesystem::path>						Paths2	 = {{std::filesystem::relative(Content2[2].first, SharePath)},
													{std::filesystem::relative(Content2[3].first, SharePath)}};

	std::vector<IoBuffer> Chunks2 = WS.GetWorkspaceShareChunks(
		PathToId(RootPath),
		PathToId("second_folder"),
		std::vector<Workspaces::ChunkRequest>{{.ChunkId = PathToId(Paths2[0])}, {.ChunkId = PathToId(Paths2[1])}},
		WorkerPool);
	CHECK(Chunks2.size() == 2);
	CHECK(Chunks2[0].GetSize() == 0);
	CHECK(Chunks2[1].GetSize() == 0);

	std::optional<std::vector<Workspaces::ShareFile>> Files =
		WS.GetWorkspaceShareFiles(PathToId(RootPath), PathToId("second_folder"), true, WorkerPool);
	CHECK(Files.has_value());
	CHECK(Files.value().size() == 6);

	Chunks2 = WS.GetWorkspaceShareChunks(
		PathToId(RootPath),
		PathToId("second_folder"),
		std::vector<Workspaces::ChunkRequest>{{.ChunkId = PathToId(Paths2[0])}, {.ChunkId = PathToId(Paths2[1])}},
		WorkerPool);
	CHECK(Chunks2.size() == 2);
	CHECK(Chunks2[0].GetView().EqualBytes(Content2[2].second.GetView()));
	CHECK(Chunks2[1].GetView().EqualBytes(Content2[3].second.GetView()));

	Workspaces::ShareFile Entry =
		WS.GetWorkspaceShareChunkInfo(PathToId(RootPath), PathToId("second_folder"), PathToId(Paths2[1]), WorkerPool);
	CHECK(Entry.Id == PathToId(Paths2[1]));
	CHECK(!Entry.RelativePath.empty());
	CHECK(Entry.Size == Content2[3].second.GetSize());

	Files = WS.GetWorkspaceShareFiles(PathToId(RootPath), PathToId("second_folder"), false, WorkerPool);
	CHECK(Files.has_value());
	CHECK(Files.value().size() == 6);

	CHECK(WS.RemoveWorkspaceShare(PathToId(RootPath), PathToId("second_folder")));
	CHECK(!WS.RemoveWorkspaceShare(PathToId(RootPath), PathToId("second_folder")));

	Files = WS.GetWorkspaceShareFiles(PathToId(RootPath), PathToId("second_folder"), false, WorkerPool);
	CHECK(!Files.has_value());

	Chunks2 = WS.GetWorkspaceShareChunks(
		PathToId(RootPath),
		PathToId("second_folder"),
		std::vector<Workspaces::ChunkRequest>{{.ChunkId = PathToId(Paths2[0])}, {.ChunkId = PathToId(Paths2[1])}},
		WorkerPool);
	CHECK(Chunks2.empty());

	CHECK(WS.RemoveWorkspace(PathToId(RootPath)));
	CHECK(!WS.RemoveWorkspace(PathToId(RootPath)));
}

TEST_CASE("workspace.share.alias")
{
	using namespace std::literals;

	WorkerThreadPool WorkerPool(std::thread::hardware_concurrency());

	ScopedTemporaryDirectory								TempDir;
	std::filesystem::path									RootPath = TempDir.Path();
	std::vector<std::pair<std::filesystem::path, IoBuffer>> Content	 = GenerateFolderContent(RootPath);

	Workspaces WS;
	CHECK(WS.AddWorkspace({PathToId(RootPath), RootPath}));
	CHECK(WS.AddWorkspaceShare(PathToId(RootPath),
							   {PathToId("second_folder"), "second_folder", "my_share"},
							   [](const std::filesystem::path& Path) { return PathToId(Path); }));

	std::optional<Workspaces::ShareAlias> Alias = WS.GetShareAlias("my_share");
	CHECK(Alias.has_value());
	CHECK(!WS.GetShareAlias("my_share2").has_value());

	std::filesystem::path			   SharePath = RootPath / "second_folder";
	std::vector<std::filesystem::path> Paths	 = {{std::filesystem::relative(Content[4].first, SharePath)},
												{std::filesystem::relative(Content[6].first, SharePath)},
												{std::filesystem::relative(Content[7].first, SharePath)},
												{"the_file_that_is_not_there.txt"}};
	std::vector<IoBuffer>			   Chunks	 = WS.GetWorkspaceShareChunks(Alias->WorkspaceId,
															  Alias->ShareId,
															  std::vector<Workspaces::ChunkRequest>{{.ChunkId = PathToId(Paths[0])},
																									{.ChunkId = PathToId(Paths[1])},
																									{.ChunkId = PathToId(Paths[2])},
																									{.ChunkId = PathToId(Paths[3])}},
															  WorkerPool);
	CHECK(Chunks.size() == 4);
	CHECK(Chunks[0].GetView().EqualBytes(Content[4].second.GetView()));
	CHECK(Chunks[1].GetView().EqualBytes(Content[6].second.GetView()));
	CHECK(Chunks[2].GetView().EqualBytes(Content[7].second.GetView()));
	CHECK(Chunks[3].GetSize() == 0);

	CHECK(WS.RemoveWorkspaceShare(Alias->WorkspaceId, Alias->ShareId));

	CHECK(!WS.GetShareAlias("my_share").has_value());

	CHECK(WS.AddWorkspaceShare(PathToId(RootPath),
							   {PathToId("second_folder"), "second_folder", "my_share"},
							   [](const std::filesystem::path& Path) { return PathToId(Path); }));
	CHECK(WS.GetShareAlias("my_share").has_value());

	CHECK(WS.RemoveWorkspace(PathToId(RootPath)));
	CHECK(!WS.RemoveWorkspace(PathToId(RootPath)));

	CHECK(!WS.GetShareAlias("my_share").has_value());
}
#endif

void
workspaces_forcelink()
{
}

}  // namespace zen
